from __future__ import absolute_import



from keras import backend as K

from keras.engine import Layer

from keras.utils.generic_utils import get_custom_objects

from keras.backend import normalize_data_format



if K.backend() == 'theano':

    import theano_backend as K_BACKEND

else:

    import tensorflow_backend as K_BACKEND



class SubPixelUpscaling(Layer):

    """ Sub-pixel convolutional upscaling layer based on the paper "Real-Time Single Image

    and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network"

    (https://arxiv.org/abs/1609.05158).

    This layer requires a Convolution2D prior to it, having output filters computed according to

    the formula :

        filters = k * (scale_factor * scale_factor)

        where k = a user defined number of filters (generally larger than 32)

              scale_factor = the upscaling factor (generally 2)

    This layer performs the depth to space operation on the convolution filters, and returns a

    tensor with the size as defined below.

    # Example :

    ```python

        # A standard subpixel upscaling block

        x = Convolution2D(256, 3, 3, padding='same', activation='relu')(...)

        u = SubPixelUpscaling(scale_factor=2)(x)

        [Optional]

        x = Convolution2D(256, 3, 3, padding='same', activation='relu')(u)

    ```

        In practice, it is useful to have a second convolution layer after the

        SubPixelUpscaling layer to speed up the learning process.

        However, if you are stacking multiple SubPixelUpscaling blocks, it may increase

        the number of parameters greatly, so the Convolution layer after SubPixelUpscaling

        layer can be removed.

    # Arguments

        scale_factor: Upscaling factor.

        data_format: Can be None, 'channels_first' or 'channels_last'.

    # Input shape

        4D tensor with shape:

        `(samples, k * (scale_factor * scale_factor) channels, rows, cols)` if data_format='channels_first'

        or 4D tensor with shape:

        `(samples, rows, cols, k * (scale_factor * scale_factor) channels)` if data_format='channels_last'.

    # Output shape

        4D tensor with shape:

        `(samples, k channels, rows * scale_factor, cols * scale_factor))` if data_format='channels_first'

        or 4D tensor with shape:

        `(samples, rows * scale_factor, cols * scale_factor, k channels)` if data_format='channels_last'.

    """



    def __init__(self, scale_factor=2, data_format=None, **kwargs):

        super(SubPixelUpscaling, self).__init__(**kwargs)



        self.scale_factor = scale_factor

        self.data_format = normalize_data_format(data_format)



    def build(self, input_shape):

        pass



    def call(self, x, mask=None):

        y = K_BACKEND.depth_to_space(x, self.scale_factor, self.data_format)

        return y



    def compute_output_shape(self, input_shape):

        if self.data_format == 'channels_first':

            b, k, r, c = input_shape

            return (b, k // (self.scale_factor ** 2), r * self.scale_factor, c * self.scale_factor)

        else:

            b, r, c, k = input_shape

            return (b, r * self.scale_factor, c * self.scale_factor, k // (self.scale_factor ** 2))



    def get_config(self):

        config = {'scale_factor': self.scale_factor,

                  'data_format': self.data_format}

        base_config = super(SubPixelUpscaling, self).get_config()

        return dict(list(base_config.items()) + list(config.items()))





get_custom_objects().update({'SubPixelUpscaling': SubPixelUpscaling})
